"""
工作流引擎模块
支持多步骤工作流执行，集成CoT思维链引导
"""
import json
import logging
from typing import Dict, List, Any, Optional, Callable
from enum import Enum
from analysis_tools import TOOLS_REGISTRY, AnalysisTools

logger = logging.getLogger(__name__)


class WorkflowStepType(Enum):
    """工作流步骤类型"""
    RETRIEVE = "retrieve"  # 检索日志
    ANALYZE = "analyze"    # 分析（调用工具）
    LLM_REASON = "llm_reason"  # LLM推理
    CONDITION = "condition"  # 条件判断


class WorkflowEngine:
    """工作流引擎"""
    
    def __init__(self, log_system):
        """
        初始化工作流引擎
        
        Args:
            log_system: TopKLogSystem实例
        """
        self.log_system = log_system
        self.tools_registry = TOOLS_REGISTRY
        self.execution_history = []
    
    def build_cot_prompt(self, query: str, available_tools: Dict) -> str:
        """
        构建CoT思维链提示词，引导LLM使用工具
        
        Args:
            query: 用户查询
            available_tools: 可用工具字典
        
        Returns:
            CoT提示词
        """
        tools_description = "\n".join([
            f"{idx+1}. {name}:\n"
            f"   描述: {tool['description']}\n"
            f"   参数: {json.dumps(tool['parameters'], ensure_ascii=False, indent=6)}"
            for idx, (name, tool) in enumerate(available_tools.items())
        ])
        
        cot_prompt = f"""你是一个智能日志分析助手。你需要使用思维链（Chain-of-Thought）来分析和解决问题。

## 可用工具列表：
{tools_description}

## 思维链工作流程：
1. **理解问题（Understand）**：
   - 仔细阅读用户问题，理解分析需求
   - 确定需要什么类型的信息
   
2. **制定计划（Plan）**：
   - 思考需要调用哪些工具
   - 确定工具调用的顺序
   - 考虑工具之间的依赖关系
   
3. **执行分析（Execute）**：
   - 按照计划调用相应的工具
   - 每个工具调用后，观察结果
   - 根据结果决定下一步行动
   
4. **综合推理（Reason）**：
   - 整合所有工具返回的结果
   - 进行深度分析和推理
   - 找出问题根源和模式
   
5. **生成报告（Report）**：
   - 基于分析结果，生成清晰的分析报告
   - 提供可行的建议和解决方案

## 工具调用格式：
当需要调用工具时，请按照以下JSON格式输出：
```json
{{
    "thought": "你的思考过程，说明为什么要调用这个工具",
    "action": {{
        "tool_name": "工具名称",
        "parameters": {{
            "参数名": "参数值"
        }}
    }}
}}
```

## 完成分析标记：
当你认为已经收集足够的信息，可以生成最终分析报告时，请输出：
"FINISHED" 或 "COMPLETE" 或 "已完成"

## 工作流程：
1. 仔细分析用户问题，确定需要哪些信息
2. 根据需要逐个调用工具收集信息
3. 每次工具调用后，评估是否还需要更多信息
4. 当信息足够时，明确表示"FINISHED"或"COMPLETE"，然后系统会生成最终报告
5. 如果问题简单，也可以直接分析而不调用工具

## 当前用户问题：
{query}

请开始你的分析，使用思维链逐步思考。根据问题复杂度，决定需要调用哪些工具，并在收集足够信息后明确表示完成。
"""
        return cot_prompt
    
    def parse_llm_tool_call(self, llm_response: str) -> Optional[Dict]:
        """
        解析LLM的工具调用响应，支持处理不完整的JSON
        
        Args:
            llm_response: LLM的响应文本
        
        Returns:
            解析后的工具调用信息，如果无法解析则返回None
        """
        try:
            json_str = None
            
            # 尝试提取JSON部分
            if "```json" in llm_response:
                json_start = llm_response.find("```json") + 7
                json_end = llm_response.find("```", json_start)
                if json_end == -1:
                    # 如果没有找到结束标记，尝试找到最后一个}
                    json_end = len(llm_response)
                json_str = llm_response[json_start:json_end].strip()
            elif "{" in llm_response and "tool_name" in llm_response:
                # 尝试找到第一个 { 到最后一个 }
                json_start = llm_response.find("{")
                json_end = llm_response.rfind("}") + 1
                if json_end <= json_start:
                    # 如果JSON不完整，尝试修复
                    json_str = llm_response[json_start:]
                    # 尝试补全缺失的括号
                    open_braces = json_str.count('{')
                    close_braces = json_str.count('}')
                    if open_braces > close_braces:
                        json_str += '}' * (open_braces - close_braces)
                        # 尝试补全字符串（如果最后一个字符串未闭合）
                        if json_str.count('"') % 2 != 0:
                            last_quote = json_str.rfind('"')
                            if last_quote > json_str.rfind(':'):
                                json_str = json_str[:last_quote + 1] + '"}'
                    else:
                        json_str = json_str[:json_end] if json_end > json_start else json_str
                else:
                    json_str = llm_response[json_start:json_end]
            else:
                return None
            
            if not json_str:
                return None
            
            # 尝试解析JSON
            parsed = json.loads(json_str)
            
            # 支持两种格式
            if "action" in parsed:
                action = parsed["action"]
                if not isinstance(action, dict) or "tool_name" not in action:
                    return None
                return {
                    "thought": parsed.get("thought", ""),
                    "tool_name": action.get("tool_name"),
                    "parameters": action.get("parameters", {})
                }
            elif "tool_name" in parsed:
                return {
                    "thought": parsed.get("thought", ""),
                    "tool_name": parsed.get("tool_name"),
                    "parameters": parsed.get("parameters", {})
                }
        except json.JSONDecodeError as e:
            # JSON解析错误，尝试更宽松的方式
            logger.warning(f"JSON解析失败: {e}, 尝试修复...")
            
            # 尝试从响应中提取关键信息
            if "tool_name" in llm_response:
                # 尝试使用正则表达式提取
                import re
                tool_name_match = re.search(r'"tool_name"\s*:\s*"([^"]+)"', llm_response)
                if tool_name_match:
                    tool_name = tool_name_match.group(1)
                    # 尝试提取parameters（简化版）
                    params_match = re.search(r'"parameters"\s*:\s*(\{[^}]*\})', llm_response)
                    parameters = {}
                    if params_match:
                        try:
                            parameters = json.loads(params_match.group(1))
                        except:
                            pass
                    
                    return {
                        "thought": "",
                        "tool_name": tool_name,
                        "parameters": parameters
                    }
            
            logger.warning(f"无法解析工具调用，响应内容: {llm_response[:300]}")
            return None
        except Exception as e:
            logger.warning(f"解析工具调用失败: {e}, 响应内容: {llm_response[:200]}")
            return None
        
        return None
    
    def execute_tool(self, tool_name: str, parameters: Dict) -> Any:
        """
        执行工具调用
        
        Args:
            tool_name: 工具名称
            parameters: 工具参数
        
        Returns:
            工具执行结果
        """
        if tool_name not in self.tools_registry:
            return {"error": f"未知工具: {tool_name}"}
        
        tool_info = self.tools_registry[tool_name]
        tool_func = tool_info["function"]
        
        # 获取工具注册表中定义的参数列表
        valid_parameters = tool_info.get("parameters", {})
        
        # 过滤参数：只保留在工具注册表中定义的参数
        filtered_parameters = {}
        for param_name, param_value in parameters.items():
            if param_name in valid_parameters:
                filtered_parameters[param_name] = param_value
            else:
                logger.warning(f"工具 {tool_name} 忽略了无效参数: {param_name}={param_value}")
        
        try:
            # 执行工具函数，只传递有效参数
            result = tool_func(**filtered_parameters)
            logger.info(f"工具 {tool_name} 执行成功")
            return result
        except Exception as e:
            logger.error(f"工具 {tool_name} 执行失败: {e}")
            return {"error": str(e)}
    
    def run_workflow(self, query: str) -> Dict[str, Any]:
        """
        运行工作流
        
        Args:
            query: 用户查询
        
        Returns:
            工作流执行结果
        """
        self.execution_history = []
        
        # 第一步：检索相关日志
        logger.info(f"开始检索日志: {query}")
        retrieved_logs = self.log_system.retrieve_logs(query, top_k=20)
        
        if not retrieved_logs:
            return {
                "response": "未找到相关日志",
                "workflow_steps": self.execution_history,
                "retrieved_logs_count": 0
            }
        
        self.execution_history.append({
            "step": "retrieve",
            "result": f"检索到 {len(retrieved_logs)} 条相关日志"
        })
        
        # 构建CoT提示词
        cot_prompt = self.build_cot_prompt(query, self.tools_registry)
        
        # 构建包含检索结果的上下文
        log_context = "## 检索到的相关日志:\n"
        for i, log in enumerate(retrieved_logs[:10], 1):  # 最多显示10条
            log_context += f"日志 {i} (相关性: {log.get('score', 0):.3f}): {log.get('content', '')[:300]}...\n"
        
        full_prompt = f"{cot_prompt}\n\n{log_context}"
        
        # 执行工具调用循环（LLM自主决定何时停止）
        tool_results = []
        analysis_context = {"logs": retrieved_logs}
        MAX_SAFE_ITERATIONS = 10  # 安全上限，防止无限循环
        iteration = 0
        
        while iteration < MAX_SAFE_ITERATIONS:
            iteration += 1
            logger.info(f"第 {iteration} 次工具调用")
            
            # 调用LLM，让它决定使用什么工具或是否完成分析
            # OllamaLLM的invoke方法接受字符串，所以需要构建完整的prompt
            user_prompt = f"""基于当前上下文，请继续分析。当前已执行的工具调用：{len(tool_results)}次。

请判断：
1. 如果需要继续分析，请调用相应的工具（输出JSON格式）
2. 如果已经收集足够的信息，可以生成最终答案，请输出 "FINISHED" 或 "COMPLETE"

请做出决定："""
            combined_prompt = f"系统提示：{full_prompt}\n\n用户问题：{user_prompt}"
            llm_response = self.log_system.llm.invoke(combined_prompt)
            
            # 检查LLM是否表示已完成
            if "FINISHED" in llm_response.upper() or "COMPLETE" in llm_response.upper() or "完成" in llm_response:
                logger.info("LLM表示已完成分析，准备生成最终答案")
                break
            
            # 解析工具调用
            tool_call = self.parse_llm_tool_call(llm_response)
            
            if not tool_call:
                # LLM没有调用工具，可能是要直接回答
                logger.info("LLM未调用工具，准备生成最终答案")
                break
            
            # 准备工具参数
            tool_params = tool_call["parameters"].copy()
            
            # 特殊处理：如果参数是"logs"，使用检索到的日志
            if "logs" in tool_params:
                # 如果值是字符串"retrieved_logs"，替换为实际日志
                if isinstance(tool_params["logs"], str) and tool_params["logs"] == "retrieved_logs":
                    tool_params["logs"] = retrieved_logs
            
            # 确保logs参数存在（大多数工具需要）
            if tool_params.get("logs") is None and "logs" in self.tools_registry[tool_call["tool_name"]]["parameters"]:
                tool_params["logs"] = retrieved_logs
            
            # 执行工具
            tool_result = self.execute_tool(tool_call["tool_name"], tool_params)
            tool_results.append({
                "iteration": iteration,
                "tool_name": tool_call["tool_name"],
                "thought": tool_call.get("thought", ""),
                "result": tool_result
            })
            
            self.execution_history.append({
                "step": f"tool_call_{iteration}",
                "tool": tool_call["tool_name"],
                "thought": tool_call.get("thought", ""),
                "result_summary": str(tool_result)[:200]
            })
            
            # 更新提示词，包含工具执行结果
            full_prompt += f"\n\n## 工具调用结果 {iteration}:\n"
            full_prompt += f"工具: {tool_call['tool_name']}\n"
            full_prompt += f"思考: {tool_call.get('thought', '')}\n"
            full_prompt += f"结果: {json.dumps(tool_result, ensure_ascii=False, indent=2)}\n"
            full_prompt += "\n请基于这个结果，判断是否已经收集足够的信息进行分析，或者是否需要继续调用其他工具。如果信息已足够，可以输出 'FINISHED' 表示完成。"
        
        # 检查是否达到安全上限
        if iteration >= MAX_SAFE_ITERATIONS:
            logger.warning(f"达到最大工具调用次数上限 ({MAX_SAFE_ITERATIONS})，自动停止")
        
        # 生成最终分析报告
        final_prompt = f"""
基于以上所有工具调用结果，请生成一份完整的分析报告。

## 用户原始问题：
{query}

## 工具调用历史：
{json.dumps(tool_results, ensure_ascii=False, indent=2)}

## 请生成：
1. 问题分析：基于工具结果，深入分析用户问题的根本原因
2. 关键发现：列出最重要的发现和模式
3. 建议方案：提供可行的解决方案或改进建议
4. 总结：简洁总结分析结论

请以结构化、专业的方式呈现报告。
"""
        
        # OllamaLLM的invoke方法接受字符串
        final_system_prompt = "你是一个专业的日志分析专家，擅长生成清晰的分析报告。\n\n"
        final_response = self.log_system.llm.invoke(final_system_prompt + final_prompt)
        
        return {
            "response": final_response,
            "workflow_steps": self.execution_history,
            "tool_calls": tool_results,
            "retrieved_logs_count": len(retrieved_logs),
            "analysis_summary": {
                "tools_used": len(tool_results),
                "tools_list": [r["tool_name"] for r in tool_results]
            }
        }
